// Copyright (c) 2018-2023 The MobileCoin Foundation

//! Utilities for mobilecoind unit tests

pub use mc_blockchain_types::BlockVersion;
pub use mc_ledger_db::test_utils::{add_block_to_ledger, add_txos_to_ledger};

use crate::{
    database::Database,
    monitor_store::{MonitorData, MonitorId},
    payments::TransactionsManager,
    service::Service,
};
use grpcio::{ChannelBuilder, EnvBuilder};
use mc_account_keys::{AccountKey, PublicAddress, DEFAULT_SUBADDRESS_INDEX};
use mc_common::logger::{log, Logger};
use mc_connection::{Connection, ConnectionManager};
use mc_connection_test_utils::{test_client_uri, MockBlockchainConnection};
use mc_consensus_scp::QuorumSet;
use mc_crypto_keys::{RistrettoPrivate, RistrettoPublic};
use mc_fog_report_validation_test_utils::{FogPubkeyResolver, MockFogResolver};
use mc_ledger_db::{test_utils::recreate_ledger_db, Ledger, LedgerDB};
use mc_ledger_sync::PollingNetworkState;
use mc_mobilecoind_api::{mobilecoind_api::MobilecoindApiClient, MobilecoindUri};
use mc_rand::{CryptoRng, RngCore};
use mc_transaction_core::{
    encrypted_fog_hint::EncryptedFogHint, onetime_keys::create_shared_secret,
    ring_signature::KeyImage, tokens::Mob, tx::TxOut, Amount, FeeMap, Token, TokenId,
};
use mc_util_from_random::FromRandom;
use mc_util_grpc::ConnectionUriGrpcioChannel;
use mc_util_uri::{ConnectionUri, FogUri};
use mc_watcher::watcher_db::WatcherDB;
use std::{
    str::FromStr,
    sync::{Arc, RwLock},
};
use tempfile::TempDir;

/// The amount each recipient gets in the test ledger.
pub const DEFAULT_PER_RECIPIENT_AMOUNT: u64 = 5_000 * 1_000_000_000_000;

/// Number of initial blocks generated by `get_testing_environment`;
pub const GET_TESTING_ENVIRONMENT_NUM_BLOCKS: usize = 10;

/// Sets up ledger_db and mobilecoind_db. Each block will contains one txo per
/// recipient.
///
/// # Arguments
/// *
/// * `num_random_recipients` - Number of random recipients to create.
/// * `known_recipients` - A list of known recipients to create.
/// * `num_blocks` - Number of blocks to create in the ledger_db.
/// * `logger`
/// * `rng`
///
/// Note that all txos will be controlled by the subindex
/// DEFAULT_SUBADDRESS_INDEX
pub fn get_test_databases(
    block_version: BlockVersion,
    num_random_recipients: u32,
    known_recipients: &[PublicAddress],
    num_blocks: usize,
    logger: Logger,
    mut rng: &mut (impl CryptoRng + RngCore),
) -> (LedgerDB, Database) {
    let mut public_addresses: Vec<PublicAddress> = (0..num_random_recipients)
        .map(|_i| mc_account_keys::AccountKey::random(&mut rng).default_subaddress())
        .collect();

    public_addresses.extend(known_recipients.iter().cloned());

    // Note that TempDir manages uniqueness by constructing paths
    // like: /tmp/ledger_db.tvF0XHTKsilx
    let ledger_db_tmp = TempDir::new().expect("Could not make tempdir for ledger db");
    let ledger_db_path = ledger_db_tmp.path();
    let mut ledger_db = recreate_ledger_db(ledger_db_path);

    for block_index in 0..num_blocks {
        let (block_version, key_images) = if block_index == 0 {
            (BlockVersion::ZERO, vec![])
        } else {
            (block_version, vec![KeyImage::from(rng.next_u64())])
        };
        add_block_to_ledger(
            &mut ledger_db,
            block_version,
            &public_addresses,
            Amount::new(DEFAULT_PER_RECIPIENT_AMOUNT, Mob::ID),
            &key_images,
            rng,
        )
        .unwrap();
    }

    let mobilecoind_db_tmp = TempDir::new().expect("Could not make tempdir for mobilecoind db");
    let mobilecoind_db_path = mobilecoind_db_tmp.path();
    let mobilecoind_db =
        Database::new(mobilecoind_db_path, logger).expect("failed creating new mobilecoind db");

    (ledger_db, mobilecoind_db)
}

pub fn get_test_monitor_data_and_id(
    rng: &mut (impl CryptoRng + RngCore),
) -> (MonitorData, MonitorId) {
    let account_key = AccountKey::random(rng);

    let data = MonitorData::new(
        account_key,
        DEFAULT_SUBADDRESS_INDEX, // first_subaddress
        1,                        // num_subaddresses
        0,                        // first_block
        "",                       // name
    )
    .unwrap();

    let monitor_id = MonitorId::from(&data);
    (data, monitor_id)
}

pub fn get_free_port() -> u16 {
    portpicker::pick_unused_port().expect("pick_unused_port")
}

pub fn get_test_fee_map() -> FeeMap {
    FeeMap::try_from_iter([
        (Mob::ID, Mob::MINIMUM_FEE),
        (TokenId::from(1), 20_480),
        (TokenId::from(2), 1_024_000),
    ])
    .unwrap()
}

pub fn setup_server<FPR: FogPubkeyResolver + Default + Send + Sync + 'static>(
    logger: Logger,
    ledger_db: LedgerDB,
    mobilecoind_db: Database,
    watcher_db: Option<WatcherDB>,
    fog_resolver_factory: Option<Arc<dyn Fn(&[FogUri]) -> Result<FPR, String> + Send + Sync>>,
    uri: &MobilecoindUri,
) -> (
    Service,
    ConnectionManager<MockBlockchainConnection<LedgerDB>>,
) {
    let fee_map = get_test_fee_map();

    let peer1 =
        MockBlockchainConnection::new(test_client_uri(1), ledger_db.clone(), 0, fee_map.clone());
    let peer2 = MockBlockchainConnection::new(test_client_uri(2), ledger_db.clone(), 0, fee_map);

    let node_ids = vec![
        peer1.uri().host_and_port_responder_id().unwrap(),
        peer2.uri().host_and_port_responder_id().unwrap(),
    ];
    let quorum_set = QuorumSet::new_with_node_ids(2, node_ids);

    let conn_manager = ConnectionManager::new(vec![peer1, peer2], logger.clone());

    let network_state = Arc::new(RwLock::new(PollingNetworkState::new(
        quorum_set,
        conn_manager.clone(),
        logger.clone(),
    )));

    {
        let mut network_state = network_state.write().unwrap();
        network_state.poll();
    }

    let transactions_manager = TransactionsManager::new(
        ledger_db.clone(),
        mobilecoind_db.clone(),
        conn_manager.clone(),
        fog_resolver_factory.unwrap_or_else(|| Arc::new(|_| Ok(FPR::default()))),
        logger.clone(),
    );

    let service = Service::new(
        ledger_db,
        mobilecoind_db,
        watcher_db,
        transactions_manager,
        network_state,
        uri,
        None,
        "unit-test".into(),
        logger,
    );

    (service, conn_manager)
}

pub fn setup_client(uri: &MobilecoindUri, logger: &Logger) -> MobilecoindApiClient {
    let env = Arc::new(
        EnvBuilder::new()
            .name_prefix("gRPC-mobilecoind-tests")
            .build(),
    );
    let ch = ChannelBuilder::new(env).connect_to_uri(uri, logger);
    MobilecoindApiClient::new(ch)
}

/// Create a ready test environment.
/// Recipients can be randomly gernerated or passed in.
/// The ledger has GET_TESTING_ENVIRONMENT_NUM_BLOCKS blocks. Each block has one
/// txo per recipient. Monitors are created for each provided MonitorData.
/// The function delays return until all monitors have processed the entire
/// ledger.
///
/// # Arguments
/// * `num_random_recipients` - random recipients to add
/// * `recipients` - particular recipient public addresses to add
/// * `monitors` - MonitorData objects specifying monitors to add
/// * `logger`
/// * `rng`
pub fn get_testing_environment(
    block_version: BlockVersion,
    num_random_recipients: u32,
    recipients: &[PublicAddress],
    monitors: &[MonitorData],
    logger: Logger,
    mut rng: &mut (impl CryptoRng + RngCore),
) -> (
    LedgerDB,
    Database,
    MobilecoindApiClient,
    Service,
    ConnectionManager<MockBlockchainConnection<LedgerDB>>,
) {
    let (ledger_db, mobilecoind_db) = get_test_databases(
        block_version,
        num_random_recipients,
        recipients,
        GET_TESTING_ENVIRONMENT_NUM_BLOCKS,
        logger.clone(),
        &mut rng,
    );
    let port = get_free_port();

    let uri =
        MobilecoindUri::from_str(&format!("insecure-mobilecoind://127.0.0.1:{port}/")).unwrap();

    log::debug!(logger, "Setting up server {:?}", port);
    let (server, server_conn_manager) = setup_server::<MockFogResolver>(
        logger.clone(),
        ledger_db.clone(),
        mobilecoind_db.clone(),
        None,
        None,
        &uri,
    );
    log::debug!(logger, "Setting up client {:?}", port);
    let client = setup_client(&uri, &logger);

    for data in monitors {
        mobilecoind_db
            .add_monitor(data)
            .expect("failed adding monitor");
    }

    wait_for_monitors(&mobilecoind_db, &ledger_db, &logger);

    (
        ledger_db,
        mobilecoind_db,
        client,
        server,
        server_conn_manager,
    )
}

/// Waits until all monitors are current with the last block of the ledger DB
///
/// # Arguments
/// * `mobilecoind_db` - Database instance
/// * `ledger_db` - LedgerDB instance
/// * `logger`
pub fn wait_for_monitors(mobilecoind_db: &Database, ledger_db: &LedgerDB, logger: &Logger) {
    let num_blocks = ledger_db.num_blocks().unwrap();

    let mut monitor_map_len: usize;
    std::thread::sleep(std::time::Duration::from_secs(1));

    'outer: loop {
        let monitor_map = mobilecoind_db
            .get_monitor_map()
            .expect("failed getting monitor map");

        monitor_map_len = monitor_map.len();

        for (i, (_monitor_id, data)) in monitor_map.iter().enumerate() {
            if data.next_block < num_blocks {
                log::info!(
                    logger,
                    "waiting for monitor {}/{}: {} of {} blocks processed",
                    i + 1, // display ordinal rather than index
                    monitor_map_len,
                    data.next_block,
                    num_blocks
                );
                std::thread::sleep(std::time::Duration::from_secs(1));
                continue 'outer;
            }
        }
        break;
    }

    if monitor_map_len > 0 {
        let plurality_char = if monitor_map_len == 1 { "" } else { "s" };
        log::info!(
            logger,
            "waited for {} monitor{} to finish processing {} blocks",
            monitor_map_len,
            plurality_char,
            num_blocks
        );
    }
}

pub fn create_tx_out(
    amount: Amount,
    recipient: &PublicAddress,
    rng: &mut (impl CryptoRng + RngCore),
) -> (TxOut, RistrettoPublic) {
    let tx_private_key = RistrettoPrivate::from_random(rng);
    let shared_secret = create_shared_secret(recipient.view_public_key(), &tx_private_key);
    let tx_out = TxOut::new(
        BlockVersion::MAX,
        amount,
        recipient,
        &tx_private_key,
        EncryptedFogHint::fake_onetime_hint(rng),
    )
    .unwrap();

    (tx_out, shared_secret)
}
